from flask import Flask, jsonify, request
import datetime, argparse, time, base64, json, traceback,queue,sys, getopt,threading
import os, cv2, datetime, requests,  socket, struct, fcntl
import numpy as np
import logging,concurrent_log
from logging import handlers
from paddlex import create_model
from PIL import Image
import numpy as np
from io import BytesIO
from datetime import datetime
import pytz,time


class Logger(object):
    level_relations = {
        'debug':logging.DEBUG,
        'info':logging.INFO,
        'warning':logging.WARNING,
        'error':logging.ERROR,
        'crit':logging.CRITICAL
    }#日志级别关系映射

    def __init__(self,filename,level='info',fmt='%(asctime)s - %(levelname)s: %(message)s'):
        path_log_f = os.path.abspath(os.path.dirname(filename))
        if not os.path.isdir(path_log_f):
            os.makedirs(path_log_f)
        
        self.logger = logging.getLogger(filename)
        self.logger.propagate = False
        self.logger.setLevel(self.level_relations.get(level))#设置日志级别
        # sh = logging.StreamHandler()#往屏幕上输出
        # th = handlers.TimedRotatingFileHandler(filename=filename, when="MIDNIGHT",backupCount=15,encoding='utf-8')#往文件里写入#指定间隔时间自动生成文件的处理器
        th = concurrent_log.ConcurrentTimedRotatingFileHandler(filename=filename, when="MIDNIGHT",backupCount=15,encoding='utf-8')#往文件里写入#指定间隔时间自动生成文件的处理器  backupCount=backCount(默认0不会删除日志)
        #实例化TimedRotatingFileHandler
        #interval是时间间隔，backupCount是备份文件的个数，如果超过这个个数，就会自动删除，when是间隔的时间单位，单位有以下几种：
        # S 秒
        # M 分
        # H 小时、
        # D 天、
        # W 每星期（interval==0时代表星期一）
        # midnight 每天凌晨
        # th.setFormatter(format_str)
        # self.logger.addHandler(sh) #把对象加到logger里
        self.logger.addHandler(th)


def get_args_parser():
    parser = argparse.ArgumentParser('online crnn en_ocr', add_help=False)
    parser.add_argument('--port', default=8080 )
    parser.add_argument('--threaded', default=False )
    parser.add_argument('--det_model_dir', default="./online_model_dir/YOLOE_P_L_1-0-0/inference" )
    parser.add_argument('--rec_model_dir', default="" )
    parser.add_argument('--save_in_log', default="./log/log_data.log" )

    return parser
parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
params = parser.parse_args([])
params.gpu_id="0"
print("初始化参数",params)

Log = Logger(params.save_in_log, level='info')  # 创建输入输出日志对象



model_det_xuanze = create_model(model_name="PP-YOLOE_plus-L",model_dir=params.det_model_dir)
# model = create_model(model_name="PP-YOLOE_plus-L",model_dir="output/dataset_xuanze_1000/269/inference",)
model_ocr_xuanze = create_model(model_name="PP-OCRv5_server_rec")




def save_img(img_src):
    # 显式使用北京时间
    tz = pytz.timezone('Asia/Shanghai')
    now = datetime.now(tz)
    timestamp_milliseconds = int(round(time.time() * 1000))  # 当前时间戳（毫秒）
    img_name=now.strftime("%Y-%m-%d_%H-%M-%S")+"_"+ str(timestamp_milliseconds)+".png"
    img_dir=os.path.join("./log",img_name.split("_")[0])
    if not os.path.isdir(img_dir):
        os.makedirs(img_dir)
    img_path=os.path.join(img_dir,img_name)
    # print(img_path)
    cv2.imwrite(img_path,img_src)
    return img_path

def cropped_img(img,coordinate):
    # 定义矩形框的坐标：左上 (x1, y1)，右下 (x2, y2)
    # 截取矩形区域（注意顺序是 [y1:y2, x1:x2]）
    coordinate[0]-=2
    coordinate[1]-=2
    coordinate[2]+=2
    coordinate[3]+=2
    return img[coordinate[1]:coordinate[3], coordinate[0]:coordinate[2]]

def output_det_xuanze(img):
    output_dict_res={}
    output_det_xuanze = model_det_xuanze.predict(img,batch_size=1,threshold=0.5)
    for det_res in output_det_xuanze:
        output_dict_res=det_res.json["res"]
        boxes=output_dict_res["boxes"]
        for boxe in boxes:
            boxe["coordinate"]=[int(i) for i in boxe["coordinate"]]
            coordinate=boxe["coordinate"]
            if boxe["cls_id"]==1:
                color=(255,0,0)
            else:
                color=(0,255,0)
                ocr_img=cropped_img(img,coordinate)
                boxe["text"]=output_ocr_xuanze(ocr_img)
            # img = cv2.rectangle(img, coordinate[:2], coordinate[2:], color, 2)
            # print(boxe)
        
        # cv2.imwrite("/paddle/PaddleX/output/train_result.png",img)
    return output_dict_res

def output_ocr_xuanze(ocr_img):
    output_ocr_xuanze = model_ocr_xuanze.predict(ocr_img, batch_size=1)
    text_re=""
    for res in output_ocr_xuanze:
        text_re=res["rec_text"]
    return text_re




# from waitress import serve
app = Flask(__name__)
# 设置最大请求体大小为 32 MB（可根据需要调整）
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 单位是字节
s = requests.session()
s.keep_alive = False

@app.route('/predict', methods=['get', 'POST'])
def add_logo():
    # json_data = request.form.to_dict()
    json_data = {}
    # print(json_data)

    try:
        # pic_url = json_data.get('pic_url')
        # base_64 = json_data.get('image_base64')
        # decoded_data = base64.b64decode(base_64) 
        # img_array = np.frombuffer(decoded_data, np.uint8)  # 转换np序列
        # base_img = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)  # 转换Opencv格 #


        # 检查是否有文件上传
        if 'file' not in request.files:
            return jsonify({'code': 400, 'info': 'No file part in the request'})

        file = request.files['file']
        # 检查是否选中了文件（即不是空文件名）
        if file.filename == '':
            return jsonify({'code': 400, 'info': 'No selected file'})
        else:
            json_data["filename"]=file.filename
        # 读取文件内容（可以进行 OCR 处理）
        image = Image.open(BytesIO(file.read()))
        image_np = np.array(image)


        # 判断是否是单通道图片
        if len(image_np.shape)==2:
            image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2BGR)

        # # 将二进制数据编码为Base64字符串
        # encoded_image = str(base64.b64encode(image_np))[2:-1]
        # json_data["img_base64"]=encoded_image
        img_path=save_img(image_np)
        json_data["img_path"]=img_path

        json_data["res"]=output_det_xuanze(image_np)
        json_data["code"]=200

    except:
        tb=traceback.format_exc()
        json_data["info"]=str(tb)
        json_data["code"]=500

    Log.logger.info(json.dumps(json_data,ensure_ascii=False))
    json_data["img_base64"]="img"
    return jsonify(json_data)


if __name__ == '__main__':
    gpu_id = 0
    # print("进入主线程：+++++++++++++++++++++++")
    app.run(host="0.0.0.0", port=params.port, threaded=params.threaded)  # ,processes=2